#!/usr/bin/env python3
import rospy
from drone_pkg.msg import QRCode
import cv2
import numpy as np
from pyzbar.pyzbar import decode
import os
import time

def Laser_Init():
    os.system("gpio mode 10 out")
    os.system("gpio mode 13 out")
    os.system("gpio write 10 1")
    os.system("gpio write 13 1")

def Laser_Control(laser_id, state):
    if laser_id == 1 and state == 0:
        os.system("gpio write 10 1")
    if laser_id == 1 and state == 1:
        os.system("gpio write 10 0")
    if laser_id == 0 and state == 0:
        os.system("gpio write 13 1")
    if laser_id == 0 and state == 1:
        os.system("gpio write 13 0")

def Laser_ON(laser_id, t):
    Laser_Control(laser_id, 1)
    time.sleep(t * 0.1)
    Laser_Control(laser_id, 0)

def shutdown():
    Laser_Control(1, 0)
    Laser_Control(0, 0)
    os.system("gpio write 19 1")
    cv2.destroyAllWindows()  

def detect_color(image, CamID):
    global pub
    # 获取图像尺寸
    height, width = image.shape[:2]

    # 在图像中心画一个红点
    center_x, center_y = width // 2, height // 2
    cv2.circle(image, (center_x, center_y), 5, (0, 0, 255), -1)

    # 转换图像为HSV色彩空间
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    # 定义黄色和绿色的HSV范围
    lower_yellow = np.array([20, 100, 100])
    upper_yellow = np.array([40, 255, 255])

    lower_green = np.array([35, 50, 50])
    upper_green = np.array([85, 255, 255])

    # 分别检测黄色和绿色
    mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)
    mask_yellow = cv2.morphologyEx(mask_yellow, cv2.MORPH_OPEN, (5, 5), iterations=5)
    mask_green = cv2.inRange(hsv, lower_green, upper_green)
    mask_green = cv2.morphologyEx(mask_green, cv2.MORPH_OPEN, (5, 5), iterations=5)

    # 查找黄色区域
    contours_yellow, _ = cv2.findContours(mask_yellow, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for contour in contours_yellow:
        # 获取矩形边界框
        x, y, w, h = cv2.boundingRect(contour)
        if w * h > 200:
            # 画矩形框
            cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 255), 2)

            # 计算矩形的中心坐标
            center_x = x + w // 2
            center_y = y + h // 2

            # 在矩形中心添加坐标
            cv2.putText(image, f"({center_x},{center_y},{w},{h})", (center_x, center_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)

            # 创建消息并发送
            msg = QRCode()
            msg.Data = '1'  # 黄色对应1
            msg.CamID = CamID
            msg.x = x
            msg.y = y
            msg.Width = w
            msg.Height = h
            pub.publish(msg)

    # 查找绿色区域
    contours_green, _ = cv2.findContours(mask_green, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for contour in contours_green:
        # 获取矩形边界框
        x, y, w, h = cv2.boundingRect(contour)
        if w * h > 200:
            # 画矩形框
            cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)

            # 计算矩形的中心坐标
            center_x = x + w // 2
            center_y = y + h // 2

            # 在矩形中心添加坐标
            cv2.putText(image, f"({center_x},{center_y},{w},{h})", (center_x, center_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

            # 创建消息并发送
            msg = QRCode()
            msg.Data = '2'  # 绿色对应2
            msg.CamID = CamID
            msg.x = x
            msg.y = y
            msg.Width = w
            msg.Height = h
            pub.publish(msg)

    # 显示图像
    cv2.imshow("Detected Colors", image)
    cv2.waitKey(1)

if __name__ == '__main__':
    global pub
    rospy.init_node("QRCode_detect", anonymous=True)
    rospy.on_shutdown(shutdown)
    os.system("gpio mode 19 out")
    os.system("gpio write 19 1")
    pub = rospy.Publisher("/QRCode_detect", QRCode, queue_size=10)
    rate = rospy.Rate(100)
    
    cap1 = cv2.VideoCapture('/dev/video0')
    
    count = 0
    flag = 0
    while not rospy.is_shutdown():
        ret, frame1 = cap1.read()
        count += 1
        try:   
            detect_color(frame1, 1)
            if (count >= 10):
                count = 0
                flag = 1 - flag
                if (flag):
                    os.system("gpio write 19 0")
                else:
                    os.system("gpio write 19 1")

        except Exception as e:
            print(f"Error: {e}")

        cv2.waitKey(1)
        rate.sleep()

    cv2.destroyAllWindows()
